#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版LSTM-CRF模型演示脚本
"""

import os
import sys
import torch
import numpy as np

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from algorithms.enhanced_lstm_crf import EnhancedLstmCRFModel

def demo_enhanced_lstm_crf():
    """演示增强版LSTM-CRF模型功能"""
    print("💎 增强版LSTM-CRF模型演示")
    print("=" * 60)
    
    # 创建测试数据
    batch_size = 2
    seq_length = 10
    input_dim = 5
    output_dim = 10
    output_seq_length = 5
    
    # 创建模型
    print(" 创建增强版LSTM-CRF模型...")
    model = EnhancedLstmCRFModel(
        input_dim=input_dim,
        hidden_dim=64,
        output_dim=output_dim,
        output_seq_length=output_seq_length,
        num_layers=2,
        dropout=0.2
    )
    
    print(f" 模型参数数量: {sum(p.numel() for p in model.parameters()):,}")
    
    # 创建测试输入
    test_input = torch.randn(batch_size, seq_length, input_dim)
    print(f"📥 测试输入形状: {test_input.shape}")
    
    # 创建序列LSTM预测结果作为增强因子
    sequence_lstm_preds = torch.randn(batch_size, 10)
    print(f"📥 序列LSTM预测形状: {sequence_lstm_preds.shape}")
    
    # 测试前向传播（训练模式）
    print("\n🔄 测试训练模式前向传播...")
    try:
        # 创建目标标签
        targets = torch.randint(0, output_dim, (batch_size, output_seq_length))
        print(f"🎯 目标标签形状: {targets.shape}")
        
        # 创建掩码
        mask = torch.ones(batch_size, output_seq_length, dtype=torch.bool)
        
        # 计算损失
        loss = model(test_input, sequence_lstm_preds, targets, mask)
        print(f"📉 训练损失: {loss.item():.4f}")
        print(" 训练模式前向传播成功")
    except Exception as e:
        print(f" 训练模式前向传播失败: {e}")
        import traceback
        traceback.print_exc()
    
    # 测试前向传播（推理模式）
    print("\n 测试推理模式前向传播...")
    try:
        model.eval()
        with torch.no_grad():
            predictions = model(test_input, sequence_lstm_preds)
            print(f" 预测结果: {predictions}")
            print(" 推理模式前向传播成功")
    except Exception as e:
        print(f" 推理模式前向传播失败: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n📋 模型架构说明:")
    print("  1. 增强版LSTM-CRF模型在传统LSTM-CRF基础上增加了序列LSTM预测融合机制")
    print("  2. 序列LSTM预测结果作为额外特征输入，增强模型的序列建模能力")
    print("  3. 模型专门为排列5设计，支持5个位置的独立预测")
    print("  4. 使用CRF层确保预测结果的序列一致性")
    print("  5. 支持GPU加速训练和推理")
    
    print("\n🏁 演示完成!")

if __name__ == "__main__":
    demo_enhanced_lstm_crf()